"""Burgers equation with source term"""

# PyHCL imports
from hogs.grids import grid1d as grid
import hogs.solvers.burgers1d as solvers
import hogs.solvers.flux.burgers.flux as flux
from hogs.solvers.time_step_functions import BurgersTimeStep

import numpy

def sfunc(x, q):
    return q[0] * 3 * (2*x - 1)

# initialize the grid
g = grid.Grid1D()
g.initialize(xlow=0, xhigh=1, dx=0.01, nb=2, nvar=1)

# create the solver
solver = solvers.BurgersWithSource1D(tf=2, pfreq=10, sfunc=sfunc, grid=g)

# add the flux function
solver.flux_function = flux.Godunov()

# set the grid for the flux function
solver.flux_function.set_grid( solver.grid )

# time step function
solver.time_step_function = BurgersTimeStep(solver.grid)

# process command line
solver.setup()

q = g.q; x = g.xc

for i, j in enumerate(x):
    if j < 0.15:
        q[0][i] = 1.0
    else:
        q[0][i] = -0.1

solver.solve()
